/**
 * @file graph_edit_001.cpp
 *
 * Copyright (c) Huawei Technologies Co., Ltd. 2012-2018. All rights reserved.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
 */

#include <iostream>
#include <fstream>
#include <algorithm>
#include <string.h>
#include "graph_utils.h"
#include "utils.h"

using namespace std;
using namespace ge;

const char *g_nodeAbs = "abs";
const char *g_nodeAdd = "add";
const char *g_nodeData = "data";

bool CreateGraph(GraphPtr &graph, AscendString graph_name)
{
    auto shape_data = vector<int64_t>({ 1, 1, 28, 28 });
    TensorDesc desc_data(ge::Shape(shape_data), FORMAT_ND, DT_FLOAT16);

    // data op
    auto data = op::Data(g_nodeData);
    data.update_input_desc_x(desc_data);
    data.update_output_desc_y(desc_data);
    auto abs = op::Abs(g_nodeAbs).set_input_x(data);
    auto add = op::Add(g_nodeAdd).set_input_x1(data).set_input_x2(abs);

    std::vector<Operator> inputs { data };
    std::vector<Operator> targets { add };
    
    graph =  Graph::ConstructFromInputs(inputs, graph_name);
    graph->SetTargets(targets);
    return true;
}

bool ModifyGraph(GraphPtr &graph)
{
    INFO_LOG("Modify Graph Start");
    GNode abs_node;
    GNode add_node;
    GNode data_node;
    std::vector<GNode> nodes = graph->GetAllNodes();
    graphStatus ret = GRAPH_FAILED;
    for (auto &node : nodes) {
        ge::AscendString name;
        ret = node.GetName(name);
        if (ret != GRAPH_SUCCESS) {
            ERROR_LOG("node get name Failed");
            return false;
        }
        std::string node_name(name.GetString());
        if (node_name == g_nodeAbs) {
            abs_node = node;
            INFO_LOG("Find src node: abs");
        } else if (node_name == g_nodeAdd) {
            add_node = node;
            INFO_LOG("Find dst node: add");
        } else if (node_name == g_nodeData) {
            data_node = node;
            INFO_LOG("Find dst node: data");
        }
    }
    bool contianSubgraph = false;
    // Remove node
    ret = graph->RemoveNode(abs_node, contianSubgraph);
    if (ret != GRAPH_SUCCESS) {
        ERROR_LOG("Modify graph: remove abs node failed");
        return false;
    }
    // Add data edge
    ret = graph->AddDataEdge(data_node, 0, add_node, 1);
    if (ret != GRAPH_SUCCESS) {
        ERROR_LOG("Modify graph: add data edge between data and add failed");
        return false;
    } else {
        INFO_LOG("Modify graph: add data edge between data and add success");
    }

    return true;
}

int main(int argc, char *argv[])
{
    // Genetate graph
    GraphPtr graph;
    bool ret;
    AscendString graph_name("IrGraph");

    ret = CreateGraph(graph, graph_name);
    if (!ret) {
        ERROR_LOG("Create Graph and set input output Failed");
        return OPENCANN_FAILED;
    }
    INFO_LOG("Create Graph and set input output Success");
    
    // Modify graph
    ret = ModifyGraph(graph);
    if (!ret) {
        ERROR_LOG("Remove node Failed");
        return OPENCANN_FAILED;
    } else {
        INFO_LOG("Remove node Success");
    }
 
    return OPENCANN_SUCCESS;
}
